import torch
from torch import nn
from torchvision.datasets import ImageFolder  # 直接加载数据目录，进行分类
from torchvision.transforms import transforms  # 将图像进行一定处理的工具
from torch.utils.data import DataLoader, random_split

# ==================== 图像数据集处理(数据增强、分割数据集) ===================
transform = transforms.Compose([  # 数据增强：如果模型的准确率无法上升，可以考虑数据增强
    transforms.Resize(28),  # 可以将单张图像进行缩放
    transforms.RandomRotation((-15, 15)),  # 将单张图像做-15°~15°的随机调整
    # 如果需要在图像中做其他的数据增强，可以查阅相关资料（颜色防抖动...）
    transforms.ToTensor(),
    transforms.Normalize([0.5, 0.5, 0.5], [0.5, 0.5, 0.5])  # 将图像像素归一化（提高训练效率）
])

dataset = ImageFolder("./data/numbers", transform=transform)
# 将原本的数据集分为 train(训练集) valid(验证集)
len_total = len(dataset)  # 60000
train_len = int(len_total * 0.8)  # 48000
valid_len = len_total - train_len  # 12000
train_dataset, valid_dataset = random_split(dataset, [train_len, valid_len])

train_loader = DataLoader(train_dataset, batch_size=1000, shuffle=True)
valid_loader = DataLoader(valid_dataset, batch_size=1000)

# ==================== 搭建训练模型（卷积神经网络） ===================
# 训练手写数字数据集 （28，28，3） --> (6-8,6-8,n) 最终卷积之后的图像大小必须在 6-8 之间（重点）
model = nn.Sequential(
    nn.Conv2d(3, 16, 3, 1, 1),  # (28x28x16)
    nn.ReLU(),
    nn.MaxPool2d(2, 2),  # (14x14x16)
    nn.Conv2d(16, 32, 3, 1, 1),  # (14x14x32)
    nn.ReLU(),
    nn.MaxPool2d(2, 2),  # (7x7x32)
    nn.Flatten(),  # (1568,) 将上述的图像矩阵进行拉伸（展平）
    nn.Dropout(),  # 随机削减50% 的神经元数量（训练提速）
    nn.Linear(7 * 7 * 32, 1024),
    nn.ReLU(),
    nn.Linear(1024, 10),  # 10是分类的数量
    nn.LogSoftmax(dim=-1)  # 多分类激活函数
)

# ==================== 将模型设置为显卡模型 ===================
device = torch.device("cuda")  # GPU下训练数据
# device = torch.device("cpu") # CPU下训练数据

model = model.to(device)  # 设置模型用GPU模式训练
# ==================== 开始训练 ===================
# 损失函数
criterion = nn.CrossEntropyLoss()
# 优化器
optimizer = torch.optim.SGD(model.parameters(), lr=0.01, momentum=0.9)
# 模型训练次数
epochs = 10
# 最高的准确率
best_acc = 0
# 开始训练
for epoch in range(epochs):
    losses = []  # 记录一个epoch下每一个batch_size的损失
    # 通过训练集计算loss损失
    model.train()  # 开启训练模式
    for features, labels in train_loader:
        # =========== 设置训练数据为GPU格式 ===========
        features = features.to(device)
        labels = labels.to(device)
        # ===========================================
        optimizer.zero_grad()
        # 预测
        label_predict = model(features)
        # 计算损失
        loss = criterion(label_predict, labels.long())  # 此处labels需要变为整数
        loss.backward()  # 反向传播
        losses.append(loss.item())
        # 更新w、b的参数
        optimizer.step()

    # 通过验证集计算accuracy准确率
    accs = []  # 纪录每一个批次的准确率
    model.eval()  # 开启验证模式（用一些真实的数据，测试训练的结果）
    for features, labels in valid_loader:
        # =========== 设置训练数据为GPU格式 ===========
        features = features.to(device)
        labels = labels.to(device)
        # ===========================================
        # 预测
        label_predict = model(features)
        # 返回预测概率最高的下标（分类序列号）
        predict = torch.argmax(label_predict, dim=-1)

        # 准确率
        acc = sum(predict == labels) / len(labels)
        accs.append(acc)

    # 计算平均损失
    avg_loss = sum(losses) / len(losses)
    # 计算平均准确率
    avg_acc = sum(accs) / len(accs)
    print(f"{epoch + 1} / {epochs} -- loss:{avg_loss:.4f} -- acc:{avg_acc * 100:.2f}%")  # :.4f保留四位小数

    # ==================== 保存模型 ===================
    if best_acc < avg_acc:
        # 将训练好的模型参数进行保存
        torch.save(model.state_dict(), "./save/best.pt")
        best_acc = avg_acc
